package com.robotic.websocket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Scope;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@ServerEndpoint(value = "/websocket/{userId}")
@Component
@Scope("prototype")
@EnableScheduling
public class WebSocketServerKey {

    private static final Logger logger = LoggerFactory.getLogger(WebSocketServerKey.class);

    /**
     * key: userId
     * value: 该用户的所有 WebSocket 会话
     */
    private static final Map<String, CopyOnWriteArraySet<Session>> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 所有连接的会话（如果有需要全员广播可用）
     */
    private static final Set<Session> SESSIONS = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        SESSIONS.add(session);
        SESSION_POOL.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>()).add(session);
        logger.info("用户 {} 已连接，当前该用户连接数：{}，总连接数：{}",
                userId, SESSION_POOL.get(userId).size(), SESSIONS.size());
    }

    @OnClose
    public void onClose(Session session, @PathParam("userId") String userId) {
        SESSIONS.remove(session);
        CopyOnWriteArraySet<Session> userSessions = SESSION_POOL.get(userId);
        if (userSessions != null) {
            userSessions.remove(session);
            if (userSessions.isEmpty()) {
                SESSION_POOL.remove(userId);
            }
        }
        logger.info("用户 {} 断开连接，剩余该用户连接数：{}，总连接数：{}",
                userId, userSessions != null ? userSessions.size() : 0, SESSIONS.size());
    }

    @OnMessage
    public void onMessage(String message, @PathParam("userId") String userId) {
        logger.info("来自用户 {} 的消息: {}", userId, message);
        // 这里可以做消息分发逻辑
    }

    @OnError
    public void onError(Session session, Throwable error) {
        logger.error("WebSocket发生错误", error);
    }

    /**
     * 向某个用户的所有连接发送消息
     */
    public synchronized void sendOneMessage(String userId, String message) {
        CopyOnWriteArraySet<Session> userSessions = SESSION_POOL.get(userId);
        if (userSessions != null) {
            for (Session session : userSessions) {
                if (session.isOpen()) {
                    session.getAsyncRemote().sendText(message);
                }
            }
        }
    }

    /**
     * 向所有连接发送消息（广播）
     */
    public void broadcast(String message) {
        for (Session session : SESSIONS) {
            if (session.isOpen()) {
                session.getAsyncRemote().sendText(message);
            }
        }
    }
}

